from torch.utils.data import Dataset, DataLoader
from datasets import Dataset as HFDataset  # 重命名以避免冲突

class MyDataset(Dataset):
    def __init__(self, split):
        # 根据 split 确定 .arrow 文件路径
        if split == "train":
            data_path = "./data/train/chn_senti_corp-train.arrow"
        elif split == "test":
            data_path = "./data/test/chn_senti_corp-test.arrow"
        elif split == "validation":
            data_path = "./data/validation/chn_senti_corp-validation.arrow"
        else:
            raise ValueError("无效的数据集类型，必须是 train/test/validation")

        # 加载 Hugging Face Dataset
        self.hf_dataset = HFDataset.from_file(data_path)

    def __len__(self):
        return len(self.hf_dataset)

    def __getitem__(self, idx):
        text = self.hf_dataset[idx]["text"]
        label = self.hf_dataset[idx]["label"]
        return text, label


if __name__ == "__main__":
    # 创建 PyTorch Dataset
    dataset = MyDataset("train")

    # 使用 DataLoader 进行迭代
    dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

    for data in dataloader:
        print(data)